package 左哥算法.ch06二叉树;

import org.junit.Test;

import java.util.*;

/**
 * B+树（非传统B+树）
 */
public class BPlus树 {

    static class BPlusTree {

        int degree;
        /*链表的哨兵节点*/
        Node head = new Node(null);
        Node tail = new Node(null);
        Node root = new Node(null, head, tail);

        public BPlusTree(int degree) {
            this.degree = degree;
        }

        public void add(int k) {
            Node node = search(root, k);
            int i = 0;
            for (; i < node.elements.size(); i++) {
                Element e = node.elements.get(i);
                if (k == e.v) {
                    e.cnt++;
                    return;
                }
                if (k < e.v) {
                    break;
                }
            }
            node.elements.add(i, new Element(k, 1));
            if (node.elements.size() > degree) {
                split(node);
            } else if (node.elements.size() - 1 == i) {
                /*修复父节点值*/
                repairParent(node);
            }
        }

        private void repairParent(Node cur) {
            if (cur == root) return;
            Node p = cur.parent;
            int i = 0;
            for (; i < p.children.size(); i++) {
                if (cur == p.children.get(i)) break;
            }
            Element last = cur.elements.get(cur.elements.size() - 1);
            if (p.elements.get(i) != last) {
                p.elements.set(i, last);
                repairParent(p);
            }
        }

        private void split(Node cur) {
            int m = cur.elements.size() >> 1; // 中间元素的下标
            Node p = cur.parent;
            /*如果是根节点溢出，则诞生一个新的根节点*/
            if (root == cur) {
                p = new Node(null);
                cur.parent = p;
                /*预先把坑位占着，方便后面统一操作*/
                p.elements.add(null);
                p.children.add(cur);
                root = p;
            }
            Node left = new Node(p);
            Node right = new Node(p);

            left.elements = new ArrayList<>(cur.elements.subList(0, m));
            right.elements = new ArrayList<>(cur.elements.subList(m, cur.elements.size()));

            if (cur.children.isEmpty()) {
                /*如果是叶子节点则修改链表*/
                cur.pre.next = left;
                left.pre = cur.pre;
                left.next = right;
                right.pre = left;
                right.next = cur.next;
                cur.next.pre = right;
            } else {
                /*如果是内部节点则将孩子转移到左右节点*/
                left.children = new ArrayList<>(cur.children.subList(0, m));
                right.children = new ArrayList<>(cur.children.subList(m, cur.children.size()));
                left.children.forEach(o -> o.parent = left);
                right.children.forEach(o -> o.parent = right);
            }

            List<Node> children = p.children;
            for (int i = 0; i < children.size(); i++) {
                if (cur == children.get(i)) {
                    /*将元素提上去*/
                    p.elements.set(i, left.elements.get(left.elements.size() - 1));
                    p.elements.add(i + 1, right.elements.get(right.elements.size() - 1));
                    /*孩子接上去*/
                    children.set(i, left);
                    children.add(i + 1, right);
                    break;
                }
            }
            if (p.elements.size() > degree) split(p);
        }

        public boolean del(int k) {
            Node node = search(root, k);
            int i = 0;
            for (; i < node.elements.size(); i++) {
                if (node.elements.get(i).v == k) break;
            }
            if (i >= node.elements.size()) return false;
            if (--node.elements.get(i).cnt > 0) return true;
            node.elements.remove(i);
            int limit = (degree - 1 >> 1) + 1;
            if (node.elements.size() - 1 == i) {
                /*修复父节点值*/
                repairParent(node);
            }
            if (node.elements.size() < limit) {
                slink(node);
            }
            return true;
        }

        private void slink(Node node) {
            if (node == root) return;
            int limit = (degree - 1 >> 1) + 1;
            Node p = node.parent;
            List<Node> children = p.children;
            int ci = 0;
            for (; ci < children.size(); ci++) {
                if (children.get(ci) == node) break;
            }
            /*无论从左侧借还是从右侧借，还是合并，都不会影响父节点的最后一个子节点的最后一个元素，所以无需向上修正*/
            if (ci != 0) {
                if (children.get(ci - 1).elements.size() > limit) {
                    takeLeft(p, ci - 1);
                } else {
                    merge(p, ci - 1);
                }
            } else {
                if (children.get(ci + 1).elements.size() > limit) {
                    takeRight(p, ci);
                } else {
                    merge(p, ci);
                }
            }
            if (p.elements.size() < limit) {
                if (p == root) {
                    /*如果根节点没有元素了,则让当前节点变为新的根节点*/
                    if (root.children.size() <= 1) {
                        /*node不一定是根节点，可能是空节点，第一个节点才是根节点*/
                        p = children.get(0);
                        p.parent = null;
                        root = p;
                    }
                } else {
                    slink(p);
                }
            }
        }

        private void merge(Node p, int li) {
            Node left = p.children.get(li);
            Node right = p.children.get(li + 1);
            if (left.children.isEmpty()) {
                left.next = right.next;
                right.next.pre = left;
            } else {
                left.children.addAll(right.children);
                right.children.forEach(o -> o.parent = left);
            }
            left.elements.addAll(right.elements);
            /*合并后的索引是右节点的索引，所以右节点不需要改变*/
            p.elements.remove(li);
            /*移除右节点，因为是合并到左节点*/
            p.children.remove(li + 1);
        }

        private void takeLeft(Node p, int li) {
            Node left = p.children.get(li);
            Node right = p.children.get(li + 1);
            /*左孩子的最后一个节点*/
            Element le = left.elements.remove(left.elements.size() - 1);
            /*将父元素放到当前的最左侧*/
            right.elements.add(0, le);
            /*修复父节点*/
            p.elements.set(li, left.elements.get(left.elements.size() - 1));
            if (!left.children.isEmpty()) {
                Node lc = left.children.remove(left.children.size() - 1);
                right.children.add(0, lc);
                lc.parent = right;
            }
        }

        /**
         * @param p：父节点
         * @param li：父元素的位置
         */
        private void takeRight(Node p, int li) {
            Node left = p.children.get(li);
            Node right = p.children.get(li + 1);

            Element re = right.elements.remove(0);
            left.elements.add(re);
            /*修复父节点*/
            p.elements.set(li, re);
            if (!right.children.isEmpty()) {
                Node rc = right.children.remove(0);
                left.children.add(rc);
                rc.parent = left;
            }
        }


        public boolean has(int k) {
            Node node = search(root, k);
            for (int i = 0; i < node.elements.size(); i++) {
                if (k == node.elements.get(i).v) return true;
                if (k < node.elements.get(i).v) break;
            }
            return false;
        }

        public long range(int l, int r) {
            Node node = search(root, l);
            long res = 0;
            for (Element e : node.elements) {
                if (e.v >= l && e.v <= r) {
                    res += (long) e.v * e.cnt;
                }
            }
            /*没找到一个符合的*/
            if (res == 0) return 0;
            node = node.next;
            while (node != tail) {
                for (Element e : node.elements) {
                    if (e.v > r) return res;
                    res += (long) e.v * e.cnt;
                }
            }
            return res;
        }

        public Node search(Node cur, int k) {
            if (cur.children.isEmpty()) return cur;
            Node next = cur.children.get(0);
            for (int i = 0; i < cur.elements.size(); i++) {
                Element e = cur.elements.get(i);
                next = cur.children.get(i);
                if (e.v >= k) break;
            }
            return search(next, k);
        }


        /**
         * B树节点，存储元素和下一个节点
         */
        public static class Node {
            Node parent;
            Node pre;
            Node next;
            List<Element> elements = new ArrayList<>();
            List<Node> children = new ArrayList<>();

            public Node(Node parent, Node pre, Node next) {
                this.parent = parent;
                this.pre = pre;
                this.next = next;
            }

            public Node(Node parent) {
                this.parent = parent;
            }

            @Override
            public String toString() {
                return elements.toString();
            }
        }

        /**
         * 每个节点中的元素
         */
        public static class Element {
            int v, cnt;

            public Element(int v, int cnt) {
                this.v = v;
                this.cnt = cnt;
            }

            @Override
            public String toString() {
                return "{" + "v=" + v + ", cnt=" + cnt + '}';
            }
        }
    }

    private static int maxn = (int) 2e4;

    static Random random = new Random();

    @Test
    public void test() {
        int[][] ops = {
                {1, 63, 1},
                {1, 522, 4},
                {1, 159, 2},
                {1, 627, 7},
                {1, 753, 8},
                {1, 608, 6},
                {1, 208, 3},
                {1, 560, 5},
        };
        BPlusTree btree = new BPlusTree(3);
        int index = 2;
        for (int i = 0; i < ops.length; i++) {
            int op = ops[i][0];
            int v = ops[i][index];
            if (op == 0) {
                System.out.println("has:" + v + "=" + btree.has(v));
            } else if (op == 1) {
                btree.add(v);
            } else if (op == 2) {
//                System.out.println("del:" + v + "=" + btree.del(v));
            } else if (op == 3) {
            }
        }
    }

    @Test
    public void test100() {
        int cnt = (int) 1e8;
        int ops = (int) 1e5;
        for (int i = 0; i < cnt; i++) {
            compare(random.nextInt(40) + 3, ops);
//            compare(3, 10);
        }
    }

    public static void print(BPlusTree.Node root) {
        Queue<BPlusTree.Node> queue = new ArrayDeque<>();
        queue.add(root);
        while (!queue.isEmpty()) {
            int size = queue.size();
            for (int i = 0; i < size; i++) {
                BPlusTree.Node poll = queue.poll();
                System.out.print(poll.elements + "\t");
                queue.addAll(poll.children);
            }
            System.out.println();
        }
    }

    @Test
    public void test2() {
        BPlusTree ans = new BPlusTree(3);
        ans.add(100);
        System.out.println(ans.range(1, 20));
    }

    public static void compare(int m, int cnt) {
        /*3位置用来做离散化*/
        int[][] ops = new int[cnt][4];
        Ans ans = new Ans();
        BPlusTree tree = new BPlusTree(m);
        try {
            for (int i = 0; i < ops.length; i++) {
                int op =random.nextBoolean() ? 1 : random.nextInt(4);
                int v = random.nextInt(maxn);
                ops[i][0] = op;
                ops[i][1] = v;
                switch (op) {
                    case 0: {
                        boolean a = ans.has(v);
                        boolean b = tree.has(v);
                        if (a != b) {
                            throw new RuntimeException("查询发生错误:预期:" + a + ";\t 结果:" + b);
                        }
                        break;
                    }
                    case 1: {
                        ans.add(v);
                        tree.add(v);
                        break;
                    }
                    case 2: {
                        /*删除一个已经存在的*/
                        v=ops[random.nextInt(Math.max(1,i))][1];
                        ops[i][1] = v;
                        boolean a = ans.del(v);
                        boolean b = tree.del(v);
                        if (a != b) {
                            throw new RuntimeException("删除发生错误:预期:" + a + ";\t 结果:" + b);
                        }
                        break;
                    }
                    case 3: {
                        int right = Math.min(v + random.nextInt(maxn), maxn);
                        ops[i][3] = right;
                        long a = ans.range(v, right);
                        long b = tree.range(v, right);
                        if (a != b) {
                            throw new RuntimeException("range发生错误:预期:" + a + ";\t 结果:" + b);
                        }
                        break;
                    }
                }
            }
        } catch (Throwable e) {
            int[] sort = new int[cnt];
            for (int j = 0; j < cnt; j++) {
                sort[j] = ops[j][1];
            }
            Arrays.sort(sort);
            for (int[] o : ops) {
                o[2] = Arrays.binarySearch(sort, o[1]) + 1;
                System.out.print("{");
                for (int i : o) {
                    System.out.print(i + ",");
                }
                System.out.println("}");
            }
            throw e;
        }
    }

    static class Ans {
        int[] cnt = new int[maxn + 10];

        public void add(int k) {
            cnt[k]++;
        }

        public boolean has(int k) {
            return cnt[k] != 0;
        }

        public boolean del(int k) {
            if (cnt[k] <= 0) {
                return false;
            }
            cnt[k]--;
            return true;
        }

        public long range(int l, int r) {
            long res = 0;
            while (l <= r) {
                res += (long) cnt[l] * l;
                l++;
            }
            return res;
        }
    }
}
